package jtg.generator;

import com.microsoft.z3.*;
import jtg.graphics.SootCFG;
import jtg.parser.Z3Parser;
import soot.Body;
import soot.Local;
import soot.Unit;
import soot.Value;
import soot.jimple.internal.JArrayRef;
import soot.jimple.internal.JAssignStmt;
import soot.jimple.internal.JIfStmt;
import soot.jimple.internal.JReturnStmt;
import soot.toolkits.graph.CompleteUnitGraph;
import soot.toolkits.graph.UnitGraph;

import java.io.File;
import java.lang.reflect.Method;
import java.math.BigInteger;
import java.util.*;

abstract public class Generator {

    protected String clsPath;
    protected String clsName;
    protected String mtdName;
    protected UnitGraph ug;
    protected Body body;

    public Generator(String className, String methodName) {
        String defaultClsPath = System.getProperty("user.dir") + File.separator + "target" + File.separator + "classes";
        new SimpleGenerator(defaultClsPath, className, methodName);
    }

    public Generator(String classPath, String className, String methodName){
        this.clsPath = classPath;
        this.clsName = className;
        this.mtdName = methodName;
        body = SootCFG.getMethodBody(clsPath, clsName, mtdName);
        SootCFG.inlineMethodBody(body);
        ug = new CompleteUnitGraph(body);
    }


    abstract void drawCFG();

    abstract List<String> generate();

    public int[] ArrayExprToJavaArray(Expr arrayExpr, int length) {
        HashMap<Integer, Object> map = new HashMap<>();
        while (arrayExpr.isStore()) {
            Expr<?> index = arrayExpr.getArgs()[1];
            Expr<?> value = arrayExpr.getArgs()[2];
            if (index.isIntNum() && value.isIntNum()) {
                map.put(((IntNum) index).getInt(), ((IntNum) value).getInt());
            }
            arrayExpr = arrayExpr.getArgs()[0];
        }
        Object other = 0;
        if (arrayExpr.isConstantArray()) {
            Expr<?> value = arrayExpr.getArgs()[0];
            if (value.isIntNum()) {
                other = ((IntNum) value).getInt();
            }
        }
        int[] result = new int[length];
        for (int i = 0; i < length; i++) {
            result[i] = (int) map.getOrDefault(i, other);
        }
        return result;
    }

    public static Object invokeMethod(String clsName, String mtdName, Object[] args) throws Exception {
        // 加载类
        Class<?> cls = Class.forName(clsName);
        Object obj = cls.newInstance();

        // 获取方法参数类型
        Class<?>[] paramTypes = new Class<?>[args.length];
        for (int i = 0; i < args.length; i++) {
            paramTypes[i] = convertToPrimitiveType(args[i].getClass());
        }

        // 获取方法
        Method method = cls.getMethod(mtdName, paramTypes);

        // 调用方法
        return method.invoke(obj, args);
    }

    private static Class<?> convertToPrimitiveType(Class<?> wrapperClass) {
        if (wrapperClass == Integer.class) return int.class;
        if (wrapperClass == Long.class) return long.class;
        if (wrapperClass == Boolean.class) return boolean.class;
        if (wrapperClass == Double.class) return double.class;
        if (wrapperClass == Float.class) return float.class;
        if (wrapperClass == Short.class) return short.class;
        if (wrapperClass == Byte.class) return byte.class;
        if (wrapperClass == Character.class) return char.class;
        return wrapperClass;
    }

    public String getTestSetByPathConstraint(Context context, List<BoolExpr> pathConstraint) {
        Solver solver = context.mkSolver();
        for (BoolExpr expr : pathConstraint) {
            solver.add(expr);
        }
        if (solver.check() == Status.SATISFIABLE) {
//            System.out.println("Path constraint is: " + pathConstraint);
//            System.out.println("The result is: " + solver.getModel());
            return getTestSetByModel(context, solver.getModel());
        } else {
            System.out.println("The result is: " + "UNSAT");
            return null;
        }
    }

    public String getTestSetByModel(Context context, Model model) {
        Z3Parser parser = new Z3Parser(context, new HashMap<>(), new HashMap<>());
        StringBuilder result = new StringBuilder();
        List<Object> paraList = new ArrayList<>();
        for (Local local : body.getParameterLocals()) {
            parser.assignLocal(local, new ArrayList<>());
            Expr paraRes = model.evaluate(parser.fromLocal(local, new ArrayList<>()), false);
            if (paraRes.isArray()) {
                int[] array = null;
                Expr length = model.evaluate(context.mkIntConst(local.getName() + "_length"), false);
                if (!length.isIntNum()) {
                    array = ArrayExprToJavaArray(paraRes, 10);
                } else {
                    array = ArrayExprToJavaArray(paraRes, ((IntNum) length).getInt());
                }
                result.append(local.getName()).append("=").append(Arrays.toString(array)).append(" ");
                paraList.add(array);
                continue;
            }
            if (parser.isFPType(local.getType())) {
                if (!(paraRes instanceof FPNum)) {
                    // 无约束
                    if (local.getType().toString().equals("double"))
                        paraList.add(0.0);
                    else
                        paraList.add(0.0f);
                    result.append(local.getName()).append("=").append(paraList.get(paraList.size() - 1)).append(" ");
                    continue;
                }
                String hex = model.evaluate(context.mkFPToIEEEBV(paraRes), false).getSExpr().substring(2);
                if (local.getType().toString().equals("double"))
                    paraList.add(Double.longBitsToDouble(new BigInteger(hex, 16).longValue()));
                else
                    paraList.add(Float.intBitsToFloat(new BigInteger(hex, 16).intValue()));
                result.append(local.getName()).append("=").append(paraList.get(paraList.size() - 1)).append(" ");
                continue;
            }
            if (paraRes.isIntNum()) {
                // 若 local 是 Enum 类型，把 int 转换为 Enum
                if (local.getType() instanceof soot.RefType) {
                    Class<?> cls = null;
                    try {
                        cls = Class.forName(((soot.RefType) local.getType()).getClassName());
                    } catch (ClassNotFoundException e) {
                        System.err.println("Class not found: " + ((soot.RefType) local.getType()).getClassName());
                        e.printStackTrace();
                    }
                    assert cls != null;
                    if (cls.isEnum()) {
                        Object enumObj = cls.getEnumConstants()[((IntNum) paraRes).getInt()];
                        paraList.add(enumObj);
                        result.append(local.getName()).append("=").append(enumObj).append(" ");
                        continue;
                    }
                }
                paraList.add(((IntNum) paraRes).getInt());
            } else if (paraRes.isBool()) {
                paraList.add(paraRes.getBoolValue());
            } else if (paraRes.isString()) {
                paraList.add(paraRes.toString().substring(1, paraRes.toString().length() - 1));
            } else {
                result.append(local.getName()).append("=").append(0).append(" ");
                paraList.add(0);
                continue;
            }
            result.append(local.getName()).append("=").append(paraRes).append(" ");
        }
        String expected = "";
        try {
            Object ret = invokeMethod(clsName, mtdName, paraList.toArray());
            expected = ret == null ? "null" : ret.toString();
        } catch (Exception e) {
            System.err.println("Error in invoking method: " + clsName + "." + mtdName);
//                e.printStackTrace();
            return result + "expected=error";
        }
        result.append("expected=").append(expected);
        return result.toString();
    }

    public List<BoolExpr> calPathConstraint(Context context, List<Unit> path) {
        Map<String, Expr> paras = new HashMap<>();
        Map<String, Integer> latestAssign = new HashMap<>();
        Z3Parser parser = new Z3Parser(context, paras, latestAssign);
        List<BoolExpr> pathConstraint = new ArrayList<>();
        for (Local local : body.getParameterLocals()) {
            latestAssign.put(local.getName(), 0);
            paras.put(local.getName() + "_0", parser.fromLocal(local, pathConstraint));
        }
        String expectedResult = "";

        for (int i = 0; i < path.size(); i++) {
            Unit stmt = path.get(i);

            if (stmt instanceof JAssignStmt) {
                Expr<?> rightOp = parser.fromSootExpr(((JAssignStmt) stmt).getRightOp());
                if (rightOp == null) {
                    // leftOp 设置为对应类型任意值
                    parser.assignLocal((Local) ((JAssignStmt) stmt).getLeftOp(), pathConstraint);
                    System.err.println("Ignore the statement: " + stmt + " assign arbitrary value to " + ((JAssignStmt) stmt).getLeftOp().toString());
                    continue;
                }
                Value leftOp = ((JAssignStmt) stmt).getLeftOp();
                if (leftOp instanceof Local) {
                    Expr<?> newLeft = parser.assignLocal((Local) leftOp, pathConstraint);
                    pathConstraint.add(parser.mkCastEq(newLeft, rightOp));
                } else if (leftOp instanceof JArrayRef) {
                    Value base = ((JArrayRef) leftOp).getBase();
                    Value index = ((JArrayRef) leftOp).getIndex();
                    if (base instanceof Local) {
                        ArrayExpr rightStore = context.mkStore(parser.fromLocal((Local) base, pathConstraint), parser.fromSootExpr(index), rightOp);
                        ArrayExpr newLeft = (ArrayExpr) parser.assignLocal((Local) base, pathConstraint);
                        pathConstraint.add(parser.mkCastEq(newLeft, rightStore));
                    } else {
                        System.out.println("Ignore the statement: " + stmt.toString());
                        continue;
                    }
                } else {
                    System.out.println("Ignore the statement: " + stmt.toString());
                    continue;
                }
                continue;
            }
            if (stmt instanceof JIfStmt) {
//                System.out.println("The ifstmt is: " + stmt);
                if (i+1 == path.size()) {
                    continue;
                }
                BoolExpr ifstms = (BoolExpr) parser.fromSootExpr(((JIfStmt) stmt).getCondition());
                Unit nextUnit = path.get(i + 1);
//                System.out.println("The nextUnit is: " + nextUnit);

                //如果ifstmt的后继语句不是ifstmt中goto语句，说明ifstmt中的条件为假
                if (!((JIfStmt) stmt).getTarget().equals(nextUnit))
                    ifstms = context.mkNot(ifstms);
                pathConstraint.add(ifstms);
                continue;
            }
            if (stmt instanceof JReturnStmt) {
                expectedResult = stmt.toString().replace("return", "").trim();
            }
        }
//        System.out.println("The step conditions with JimpleVars are: " + pathConstraint);
        return pathConstraint;
    }
}
